# examples/streaming_translation_usage.py
# 流式翻译使用示例

import asyncio
import websockets
import json
import os
import time
from pathlib import Path
import argparse

"""
这是一个流式翻译的客户端示例，演示如何使用WebSocket接口进行实时翻译。
实时翻译允许在用户说话的同时进行翻译，而不必等待整句话结束。

用法：
    python streaming_translation_usage.py --file audio.webm --token your_token
"""

async def streaming_translation(audio_file_path, token, host="localhost", port=8000):
    """使用WebSocket进行流式翻译"""
    
    # 检查音频文件是否存在
    audio_path = Path(audio_file_path)
    if not audio_path.exists():
        print(f"错误: 音频文件 {audio_file_path} 不存在")
        return
        
    # WebSocket连接URL
    uri = f"ws://{host}:{port}/api/v1/ws/realtime_translate?token={token}"
    
    # 读取音频文件
    with open(audio_path, "rb") as f:
        audio_data = f.read()
    
    # 将音频分块，模拟实时音频流
    chunk_size = 8192  # 每个块的大小（字节）
    audio_chunks = [audio_data[i:i+chunk_size] for i in range(0, len(audio_data), chunk_size)]
    
    print(f"连接到 {uri}")
    print(f"音频文件: {audio_file_path}, 大小: {len(audio_data)} 字节")
    print(f"分割为 {len(audio_chunks)} 个块")
    
    async with websockets.connect(uri) as websocket:
        # 等待连接确认
        response = await websocket.recv()
        print(f"服务器连接响应: {response}")
        
        # 发送初始化配置
        config = {
            "source_lang": "zh",  # 源语言，auto表示自动检测
            "target_lang": "en",  # 目标语言
            "chunk_duration": 2.0,  # 音频块时长（秒）
            "enable_streaming": True  # 启用流式翻译
        }
        await websocket.send(json.dumps(config))
        
        # 等待配置确认
        response = await websocket.recv()
        print(f"配置响应: {response}")
        
        # 发送音频数据块
        for i, chunk in enumerate(audio_chunks):
            print(f"发送音频块 {i+1}/{len(audio_chunks)}, 大小: {len(chunk)} 字节")
            await websocket.send(chunk)
            
            # 模拟一些处理时间
            await asyncio.sleep(0.1)
            
            # 非阻塞接收响应
            try:
                response = await asyncio.wait_for(websocket.recv(), timeout=0.01)
                print(f"流式翻译响应: {response}")
            except asyncio.TimeoutError:
                # 超时，没有收到响应，继续发送下一块
                pass
        
        # 发送停止命令
        print("发送停止命令")
        await websocket.send("stop")
        
        # 等待最终结果
        try:
            while True:
                response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                response_data = json.loads(response)
                
                # 检查是否是最终结果
                if response_data.get("type") == "final" or response_data.get("status") == "success":
                    print("\n最终翻译结果:")
                    print(f"源文本: {response_data.get('source_text')}")
                    print(f"翻译文本: {response_data.get('translated_text')}")
                    break
                else:
                    print(f"收到响应: {response}")
        except asyncio.TimeoutError:
            print("等待最终结果超时")
        except Exception as e:
            print(f"处理响应时发生错误: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="流式翻译客户端示例")
    parser.add_argument("--file", type=str, required=True, help="音频文件路径 (.webm 格式)")
    parser.add_argument("--token", type=str, required=True, help="API访问令牌")
    parser.add_argument("--host", type=str, default="localhost", help="服务器主机名")
    parser.add_argument("--port", type=int, default=8000, help="服务器端口")
    
    args = parser.parse_args()
    
    # 运行异步函数
    asyncio.run(streaming_translation(
        audio_file_path=args.file, 
        token=args.token,
        host=args.host,
        port=args.port
    ))
